{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime, timedelta\n",
    "from tqdm import tqdm\n",
    "import zipfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def get_price_data(inst, freq, start_date, end_date, fields):\n",
    "    \"\"\"\n",
    "    获取期货价格数据\n",
    "    :param inst: str, 期货代码\n",
    "    :param freq: str, 数据频率\n",
    "    :param start_date: str, 开始日期\n",
    "    :param end_date: str, 结束日期\n",
    "    :param fields: list, 需要获取的数据字段\n",
    "    :return: pd.DataFrame, 期货价格数据\n",
    "    \"\"\"\n",
    "    column_dict = {\n",
    "    \"date\":\"Date\",\n",
    "    \"datetime\":\"DateTime\",\n",
    "    \"open\":\"Open\",\n",
    "    \"high\":\"High\",\n",
    "    \"low\":\"Low\",\n",
    "    \"close\":\"Close\",\n",
    "    \"volume\":\"Volume\",\n",
    "    \"total_turnover\":\"Amount\",\n",
    "    \"open_interest\":\"OpenInterest\"\n",
    "    }\n",
    "    append_fields = [\"date\",\"datetime\"]\n",
    "    price = get_price(inst,frequency=freq, start_date=start_date, end_date=end_date,fields=fields)\n",
    "    data_columns = list(price.columns)\n",
    "    data_fields = list(set(append_fields).intersection(set(price.index.names)))\n",
    "    data_fields.extend(fields)\n",
    "\n",
    "    price = price.reset_index()\n",
    "    price = price[data_fields]\n",
    "    price = price.rename(columns=column_dict)\n",
    "    freq_num = int(freq[:-1])\n",
    "    if \"m\" in freq:\n",
    "        price['DateTime'] = price['DateTime'].apply(lambda x: x-timedelta(minutes=freq_num))\n",
    "    return price\n",
    "\n",
    "# price = get_price_data(inst, \"1d\", start_date, end_date, fields)\n",
    "# price.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_data_to_zip(zip_handler, data, freq, inst):\n",
    "    file = \"price.csv\"\n",
    "    data.to_csv(file, index=False)\n",
    "    ins_info = instruments(inst)\n",
    "    # path 期货/1d/交易所/品种.csv\n",
    "    save_path = f'data/{ins_info.type}/{freq}/{ins_info.underlying_symbol}/{ins_info.order_book_id}.csv'\n",
    "    zip_handler.write(file,arcname=save_path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_all_data_to_zip(insts,freq):\n",
    "    start_date = \"2010-01-04\"\n",
    "    end_date = (datetime.now() - timedelta(days=1)).strftime(\"%Y-%m-%d\")\n",
    "    fields = [\"open\", \"high\", \"low\", \"close\", \"volume\", \"total_turnover\", \"open_interest\"]\n",
    "    zip_handler = zipfile.ZipFile(f\"ricequant_data_{freq}.zip\", \"w\")\n",
    "    # 保存日线以下的合约 1m 5m 15m 1h\n",
    "#     inst = \"RB88\"\n",
    "    for inst in tqdm(insts):\n",
    "        try:\n",
    "            price = get_price_data(inst, freq, start_date, end_date, fields)\n",
    "            save_data_to_zip(zip_handler, price, freq, inst)\n",
    "        except Exception as e:\n",
    "            continue\n",
    "#             print(inst, \"下载出现错误\", e)\n",
    "    zip_handler.close()\n",
    "# save_hour_data_to_zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 9358/9358 [03:25<00:00, 45.64it/s]\n"
     ]
    }
   ],
   "source": [
    "# 获取所有的inst\n",
    "product = \"Future\"\n",
    " # 一小时 \"1H\" 5分钟 用5min\n",
    "freq = \"1d\" \n",
    "insts = all_instruments(type=product, market='cn', date=None)\n",
    "insts = set(insts.order_book_id)\n",
    "# insts = set(insts.underlying_symbol)\n",
    "# insts = list(insts)\n",
    "# 如果是指数的话后面加99 如果是主力合约后面加88\n",
    "# insts = [ins+\"88\" for ins in insts]\n",
    "# print(insts,len(insts))\n",
    "# save_hour_data_to_zip(insts)\n",
    "# print(\"freq:\",freq,\"insts:\",insts)\n",
    "save_all_data_to_zip(insts,freq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
